# Dickstein, Ho, and Mark (2023)

parametrize_choice_data <- function(
  theta,
  data,
  covar_names,
  covar_lens,
  score=TRUE) {
  
  if (max(grepl("^score", names(data)))){
    stop("There is a column name starting with 'score' here which is not ok. Please amend your ways.")
  }
  
  # Back out individual covar names and beta params from covar_lanes
  W1_names <- covar_names[1:covar_lens[1]]
  W2_names <- covar_names[(1+covar_lens[1]):(covar_lens[1]+covar_lens[2])]
  W3_names <- covar_names[(1+covar_lens[1]+covar_lens[2]):(covar_lens[1]+covar_lens[2]+covar_lens[3])]
  Z_names  <- covar_names[(1+covar_lens[1]+covar_lens[2]+covar_lens[3]):(covar_lens[1]+covar_lens[2]+covar_lens[3]+covar_lens[4])]
  
  beta1 <- theta[1:covar_lens[1]]
  beta2 <- theta[(1+covar_lens[1]):(covar_lens[1]+covar_lens[2])]
  beta3 <- theta[(1+covar_lens[1]+covar_lens[2]):(covar_lens[1]+covar_lens[2]+covar_lens[3])]
  beta0  <- theta[(1+covar_lens[1]+covar_lens[2]+covar_lens[3]):(covar_lens[1]+covar_lens[2]+covar_lens[3]+covar_lens[4])]
  
  # Construct covar vectors
  W1_mat <- as.matrix(data[W1_names])
  W2_mat <- as.matrix(data[W2_names])
  W3_mat <- as.matrix(data[W3_names])
  Z_mat <- as.matrix(data[Z_names])
  
  # Construct parameter vectors
  data['alpha_i'] <- exp(W1_mat %*% beta1)
  data['omega_i'] <- exp(W2_mat %*% beta2)
  data['psi_i']   <- exp(W3_mat %*% beta3)
  data['beta0_z'] <- Z_mat %*% beta0
  data['price_sens'] <- data['alpha_i'] - data['psi_i'] 
  
  # Construct mean utility
  data['delta_i'] <- data['x_av'] + (1/2)*data['omega_i']*data['x_av']^2 - (data['price_sens']*data['p']) + data['beta0_z']
  
  # Calculate predicted shares
  data <- data %>%
    mutate(exp_delta_i = exp(delta_i)) %>%
    group_by(subscriberid_year) %>%
    mutate(share_denom = sum(exp_delta_i)) %>%
    ungroup() %>%
    mutate(s_ij = exp_delta_i / share_denom) 
  if(any(data$s_ij == 0)){
    warning("There are shares that are equal to zero in the parametrized choice data.")
  }
  
  # Calculate the derivative of mean utility to various parameters
  if (score == TRUE) {
    for (i in 1:covar_lens[1]) {
      data['mean_util_deriv'] <- -(data['alpha_i']*data['p']) * data[W1_names[i]]
      data <- data %>% 
        group_by(subscriberid_year) %>%
        mutate(sum_exp_delta_interaction = sum(exp_delta_i * mean_util_deriv)) %>%
        ungroup()
      data[paste('score_W1', i, sep="_")] <- data['choice'] * 
        (data['mean_util_deriv'] - (data['sum_exp_delta_interaction'] / data['share_denom']))
    }
    for (i in 1:covar_lens[2]) {
      data['mean_util_deriv'] <- (1/2) * data['omega_i']*(data['x_av']^2) * data[W2_names[i]]
      data <- data %>% 
        group_by(subscriberid_year) %>%
        mutate(sum_exp_delta_interaction = sum(exp_delta_i * mean_util_deriv)) %>%
        ungroup()
      data[paste('score_W2', i, sep="_")] <- data['choice'] * 
        (data['mean_util_deriv'] - (data['sum_exp_delta_interaction'] / data['share_denom']))
    }
    for (i in 1:covar_lens[3]) {
      data['mean_util_deriv'] <- (data['psi_i']*data['p']) * data[W3_names[i]]
      data <- data %>% 
        group_by(subscriberid_year) %>%
        mutate(sum_exp_delta_interaction = sum(exp_delta_i * mean_util_deriv)) %>%
        ungroup()
      data[paste('score_W3', i, sep="_")] <- data['choice'] * 
        (data['mean_util_deriv'] - (data['sum_exp_delta_interaction'] / data['share_denom']))
    }
    for (i in 1:covar_lens[4]) {
      data['mean_util_deriv'] <- data[Z_names[i]]
      data <- data %>% 
        group_by(subscriberid_year) %>%
        mutate(sum_exp_delta_interaction = sum(exp_delta_i * mean_util_deriv)) %>%
        ungroup()
      data[paste('score_Z', i, sep="_")] <- data['choice'] * 
        (data['mean_util_deriv'] - (data['sum_exp_delta_interaction'] / data['share_denom']))
    }
  }
  
  return(data)
}

